from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query
from chat_manager import ConnectionManager

app = FastAPI()
manager = ConnectionManager()

@app.websocket("/ws/chat/")
async def websocket_endpoint(websocket: WebSocket, username: str = Query(...)):
    await manager.connect(username, websocket)
    try:
        while True:
            data = await websocket.receive_json()
            msg_type = data.get("type", "broadcast")
            content = data.get("message")
            to_user = data.get("to")

            if msg_type == "private" and to_user:
                await manager.send_personal_message(f"私信【{username}】→【{to_user}】：{content}", to_user)
            else:
                await manager.broadcast(f"【{username}】：{content}")
    except WebSocketDisconnect:
        await manager.disconnect(username)
        await manager.broadcast(f"用户【{username}】离开了聊天室")